import os

from PIL import Image
from torch.utils.data import DataLoader,Dataset
from torch.utils.tensorboard import SummaryWriter
from torchvision.transforms import transforms
from torchvision.utils import make_grid

"""
 读取数据使用torch,并使用Tensorboard进行可视化
"""

class MyData(Dataset):

    # 初始化 文件夹下所有文件名称列表  ； 标签 ； 根路径  # root_dir ： 根目录 ； label_dir ： 标签 也是 文件夹名称
    def __init__(self , root_dir , label_dir , transform):
        self.root_dir = root_dir
        self.label_dir = label_dir
        self.transform = transform  # PIL -> Tensor
        self.path = os.path.join(self.root_dir , self.label_dir)
        if not os.path.lexists(self.path):
            os.makedirs(self.path)
        self.imgs_list = os.listdir(self.path)

    # 返回图片和标签
    def __getitem__(self, item):
        img_name = self.imgs_list[item]
        img_path = os.path.join(self.root_dir , self.label_dir , img_name)
        img = Image.open(img_path)
        img = self.transform(img)
        label = self.label_dir
        sample = {'image': img, 'target': label}  # 返回张量图片和标签
        return sample

    def __len__(self):
        return len(self.imgs_list)


if __name__ == '__main__':
    root_dir = "../dataset/train/"
    ants_label_dir = "ants"
    bees_label_dir = "bees"
    #格式转换
    transform = transforms.Compose([transforms.Resize((256,256)) ,transforms.ToTensor()])
    ants_dataset = MyData(root_dir, ants_label_dir ,transform=transform)
    ant_img , ant_label = ants_dataset[0]
    bees_dataset = MyData(root_dir , bees_label_dir , transform=transform)
    bee_img , bee_label = bees_dataset[0]

    #训练数据集
    train_dataset = ants_dataset + bees_dataset

    dataloader = DataLoader(train_dataset , batch_size= 1 , num_workers= 2)  #数据加载器

    # 测试图片
    print(train_dataset[243]['target'])
    print(train_dataset[0]['image'].shape)


    writer = SummaryWriter("logs01")  #tensorboard

    step = 0
    for data in train_dataset:
        writer.add_image("train" , data['image']  ,global_step=step)
        print(step)
        step+=1

    for i, j in enumerate(dataloader):
        imgs, labels = j
        print(type(j))
        print(i, j['image'].shape)
        writer.add_image("train_data_b2", make_grid(j['image']), i)
    writer.close()